package thired.class09_union;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

public class Class01_UnionFind {
    public static class Node<T>{
        public T val;
        public Node(T val){
            this.val = val;
        }
    }

    public class UnionFind<T>{
        public HashMap<T,Node<T>> nodes;
        public HashMap<Node<T>,Node<T>> parent;
        public HashMap<Node<T>,Integer> size;

        public UnionFind(List<T> lists){
            nodes = new HashMap<>();
            parent = new HashMap<>();
            size = new HashMap<>();
            for (T val : lists) {
                Node<T> cur = new Node<T>(val);
                nodes.put(val,cur);
                parent.put(cur,cur);
                size.put(cur,1);
            }
        }
        public Node<T> findFather(Node<T> cur){
            Stack<Node<T>> stack = new Stack<>();
            while(parent.get(cur) != cur){
                stack.add(cur);
                cur = parent.get(cur);
            }
            while (!stack.isEmpty()){
                parent.put(stack.pop(),cur);
            }
            return cur;
        }
        public boolean isSameSet(T a, T b) {
            if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
                return false;
            }
            return findFather(nodes.get(a)) == findFather(nodes.get(b));
        }


        public void union(T a,T b){
            if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
                return;
            }
            Node<T> aHead = findFather(nodes.get(a));
            Node<T> bHead = findFather(nodes.get(b));
            if(aHead == bHead){
                return;
            }
            Node<T> big = size.get(aHead) >= size.get(bHead) ? aHead : bHead;
            Node<T> small = big == aHead ? bHead : aHead;
            parent.put(small,big);
            size.put(big,size.get(aHead) + size.get(bHead));
            size.remove(small);
        }
    }
}
